Repository navigation
Make the XLB backend an optional extra - #15
Merged
Merged
Conversation
xlb and warp-lang are only used by the XLB backend (_solve_flow_xlb.py and _warp_convergence.py), and their imports are already lazy, but as hard requirements they install XLB, JAX/jaxlib, Warp and their dependencies for every user of the default Taichi solver. Move them to a new `xlb` extra (pip install "kabs[xlb]", or uv sync --extra xlb). Selecting backend="xlb" without it raises an ImportError that names the extra. Tests: the XLB convergence tests skip without xlb, like the XLB solver tests already do; a fresh-process test checks that importing kabs loads none of xlb, jax or warp; a test checks the missing-extra error; and KABS_TEST_WITHOUT_XLB=1 checks a real install without the extra. CI runs the suite both with and without the extra.
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## dev #15 +/- ##
=======================================
+ Coverage 87.8% 88.0% +0.1%
=======================================
Files 11 11
Lines 1300 1300
=======================================
+ Hits 1142 1144 +2
+ Misses 158 156 -2 🚀 New features to boost your workflow:
|
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Why
Since the backend wrapper (27bbfeb),
xlb>=0.3.1andwarp-lang==1.10.0are in[project].dependencies, so every install of kabs pulls in XLB, JAX/jaxlib, Warp and their dependencies (ml-dtypes, opt-einsum, numpy-stl, python-utils, plus pydantic and trimesh through XLB). Only the XLB backend uses them:kabs/_solve_flow_xlb.pyandkabs/_warp_convergence.py, and their imports were already lazy.Downstream projects that only use the Taichi solver pay for that on every install. In PoreStore, which installs kabs in its development and CI environments, the XLB stack is about 340 MB of the environment on macOS arm64 (jaxlib alone is about 256 MB, Warp 65 MB) and none of it is ever imported. In kabs's own environment on Linux x86_64 the gap is larger: 1,942 MB with the XLB stack against 1,085 MB without it.
What changes
pyproject.toml:xlbandwarp-langmove to a new optional extra,xlb. Install it withpip install "kabs[xlb]", oruv sync --extra xlbin a clone. The pins are unchanged.uv.lockis regenerated with uv 0.8.17, the version CI uses; the only change is that the two packages now sit under the extra.kabs/_solve_flow_xlb.py: selecting the XLB backend without the extra (solve_flow(backend="xlb")orsolve_flow_xlb(...)) raises anImportErrorthat nameskabs[xlb]and theuv sync --extra xlbcommand, chained to the originalModuleNotFoundError. Input validation still runs before the import, so a baddirectionorcompute_backendis still reported as aValueErrorwithout XLB. The docstring's installation note points at the extra.import kabsloads none ofxlb,jaxorwarp(this was already true, and a new test now enforces it).jaxis imported only inside the JAX convergence helpers,xlbonly insidesolve_flow_xlbafter validation, andwarponly on thecompute_backend="warp"path throughkabs._warp_convergence, which nothing imports at package level.README.mdanddocs/warp-backend-readiness.md: the install instructions name the extra.Tests
tests/test_convergence_xlb.pyskips withoutxlb(pytest.importorskip), astests/test_solve_flow_xlb.pyalready did. It used to importjaxat module level, which would fail collection without the extra.tests/test_backend_import_isolation.py::test_import_kabs_does_not_import_the_xlb_stackchecks, in a fresh process, thatimport kabsdoes not loadxlb,jax,warporkabs._warp_convergence.tests/test_solve_flow_xlb_validation.py::test_missing_xlb_extra_is_named_in_the_error[jax|warp]hidesxlband checks the error through the publicsolve_flow. It runs in both environments.tests/test_solve_flow_xlb_validation.py::test_install_without_the_extra_omits_the_xlb_stackruns whenKABS_TEST_WITHOUT_XLB=1, the same opt-in pattern asKABS_TEST_WARP_CUDA. It checks thatxlb,warpandjaxare really absent and thatsolve_flow(backend="xlb")raises the named error..github/workflows/tests.ymlruns the suite twice: once withuv sync --group dev --extra xlb, which runs every test and uploads coverage as before, and once withuv sync --group devandKABS_TEST_WITHOUT_XLB=1. Tests run withuv run --no-sync, so the job's environment is exactly what its sync step installed.Results on Linux x86_64 (Python 3.12, CPU only, so the
KABS_TEST_WARP_CUDAtests skip in every run):devat 2f6d82f (before)--extra xlbKABS_TEST_WITHOUT_XLB=1uvx ruff@0.15.6 checkpasses on the changed Python files.One thing I noticed while testing, unrelated to this change: on macOS arm64,
tests/test_convergence_xlb.py::test_warp_cpu_combined_reduction_matches_masked_numpyaborts the whole pytest process (Fatal Python error: Aborted) inside Warp's library load (warp/_src/context.py,load_dll). The session fixture intests/conftest.pyhas already initialised Taichi by then; I did not check whether that is the cause. It passes on Linux, so CI is unaffected, and a macOS install without the extra skips it.